use arrow::array::{Array, ArrayRef, Float64Array};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use anyhow::{Context, Result};
use std::sync::Arc;
use crate::compute::indicators::technical_indicators::*;
use crate::compute::memory::memory_pool::MemoryPool;

pub struct FeatureEngineer {
    memory_pool: MemoryPool,
}

impl FeatureEngineer {
    pub fn new() -> Self {
        Self {
            memory_pool: MemoryPool::new(),
        }
    }

    pub fn calculate_features(&self, batch: &RecordBatch) -> Result<RecordBatch> {
        // 获取必要的列
        let close = batch
            .column_by_name("close")
            .context("Missing 'close' column")?
            .as_any()
            .downcast_ref::<Float64Array>()
            .context("'close' column is not Float64Array")?;

        let high = batch
            .column_by_name("high")
            .context("Missing 'high' column")?
            .as_any()
            .downcast_ref::<Float64Array>()
            .context("'high' column is not Float64Array")?;

        let low = batch
            .column_by_name("low")
            .context("Missing 'low' column")?
            .as_any()
            .downcast_ref::<Float64Array>()
            .context("'low' column is not Float64Array")?;

        let volume = batch
            .column_by_name("volume")
            .context("Missing 'volume' column")?
            .as_any()
            .downcast_ref::<Float64Array>()
            .context("'volume' column is not Float64Array")?;

        // 计算基础特征
        let returns = self.calculate_returns(close)?;
        let log_returns = self.calculate_log_returns(close)?;
        let volatility = self.calculate_volatility(&returns, 20)?;

        // 计算移动平均线
        let ma5 = self.calculate_ma(close, 5)?;
        let ma10 = self.calculate_ma(close, 10)?;
        let ma20 = self.calculate_ma(close, 20)?;
        let ma60 = self.calculate_ma(close, 60)?;

        // 计算指数移动平均线
        let ema5 = self.calculate_ema(close, 5)?;
        let ema10 = self.calculate_ema(close, 10)?;
        let ema20 = self.calculate_ema(close, 20)?;

        // 计算布林带
        let (bb_upper, bb_middle, bb_lower) = self.calculate_bollinger_bands(close, 20, 2.0)?;

        // 计算RSI
        let rsi = self.calculate_rsi(close, 14)?;

        // 计算成交量特征
        let volume_ma5 = self.calculate_ma(volume, 5)?;
        let volume_ma20 = self.calculate_ma(volume, 20)?;

        // 计算ATR
        let atr = self.calculate_atr(high, low, close, 14)?;

        // 创建新的Schema
        let schema = Schema::new(vec![
            Field::new("close", DataType::Float64, false),
            Field::new("returns", DataType::Float64, false),
            Field::new("log_returns", DataType::Float64, false),
            Field::new("volatility", DataType::Float64, false),
            Field::new("ma5", DataType::Float64, false),
            Field::new("ma10", DataType::Float64, false),
            Field::new("ma20", DataType::Float64, false),
            Field::new("ma60", DataType::Float64, false),
            Field::new("ema5", DataType::Float64, false),
            Field::new("ema10", DataType::Float64, false),
            Field::new("ema20", DataType::Float64, false),
            Field::new("bb_upper", DataType::Float64, false),
            Field::new("bb_middle", DataType::Float64, false),
            Field::new("bb_lower", DataType::Float64, false),
            Field::new("rsi", DataType::Float64, false),
            Field::new("volume_ma5", DataType::Float64, false),
            Field::new("volume_ma20", DataType::Float64, false),
            Field::new("atr", DataType::Float64, false),
        ]);

        // 创建新的RecordBatch
        let columns: Vec<ArrayRef> = vec![
            Arc::new(close.clone()),
            Arc::new(returns),
            Arc::new(log_returns),
            Arc::new(volatility),
            Arc::new(ma5),
            Arc::new(ma10),
            Arc::new(ma20),
            Arc::new(ma60),
            Arc::new(ema5),
            Arc::new(ema10),
            Arc::new(ema20),
            Arc::new(bb_upper),
            Arc::new(bb_middle),
            Arc::new(bb_lower),
            Arc::new(rsi),
            Arc::new(volume_ma5),
            Arc::new(volume_ma20),
            Arc::new(atr),
        ];

        Ok(RecordBatch::try_new(Arc::new(schema), columns)?)
    }

    fn calculate_returns(&self, close: &Float64Array) -> Result<Float64Array> {
        let mut returns = Vec::with_capacity(close.len());
        returns.push(0.0); // 第一个值没有收益率

        for i in 1..close.len() {
            let current = close.value(i);
            let previous = close.value(i - 1);
            let return_value = (current - previous) / previous;
            returns.push(return_value);
        }

        Ok(Float64Array::from(returns))
    }

    fn calculate_log_returns(&self, close: &Float64Array) -> Result<Float64Array> {
        let mut log_returns = Vec::with_capacity(close.len());
        log_returns.push(0.0); // 第一个值没有收益率

        for i in 1..close.len() {
            let current = close.value(i);
            let previous = close.value(i - 1);
            let log_return = (current / previous).ln();
            log_returns.push(log_return);
        }

        Ok(Float64Array::from(log_returns))
    }

    fn calculate_volatility(
        &self,
        returns: &Float64Array,
        window: usize,
    ) -> Result<Float64Array> {
        let mut volatility = Vec::with_capacity(returns.len());
        
        // 填充前window-1个值
        for _ in 0..window-1 {
            volatility.push(0.0);
        }

        for i in window-1..returns.len() {
            let window_slice: Vec<f64> = returns
                .values()
                .slice(i - window + 1, window)
                .iter()
                .copied()
                .collect();

            let mean = window_slice.iter().sum::<f64>() / window as f64;
            let variance = window_slice
                .iter()
                .map(|&x| (x - mean).powi(2))
                .sum::<f64>() / (window - 1) as f64;
            let std_dev = variance.sqrt();

            volatility.push(std_dev);
        }

        Ok(Float64Array::from(volatility))
    }

    fn calculate_ma(
        &self,
        values: &Float64Array,
        window: usize,
    ) -> Result<Float64Array> {
        let mut ma = Vec::with_capacity(values.len());
        
        // 填充前window-1个值
        for _ in 0..window-1 {
            ma.push(0.0);
        }

        for i in window-1..values.len() {
            let sum: f64 = values
                .values()
                .slice(i - window + 1, window)
                .iter()
                .sum();
            ma.push(sum / window as f64);
        }

        Ok(Float64Array::from(ma))
    }

    fn calculate_ema(
        &self,
        values: &Float64Array,
        window: usize,
    ) -> Result<Float64Array> {
        let mut ema = Vec::with_capacity(values.len());
        let alpha = 2.0 / (window as f64 + 1.0);
        
        // 第一个值使用简单平均
        let first_ma = values
            .values()
            .slice(0, window.min(values.len()))
            .iter()
            .sum::<f64>() / window.min(values.len()) as f64;
        ema.push(first_ma);

        // 计算其余值
        for i in 1..values.len() {
            let current = values.value(i);
            let previous_ema = ema[i - 1];
            let new_ema = current * alpha + previous_ema * (1.0 - alpha);
            ema.push(new_ema);
        }

        Ok(Float64Array::from(ema))
    }

    fn calculate_bollinger_bands(
        &self,
        values: &Float64Array,
        window: usize,
        num_std: f64,
    ) -> Result<(Float64Array, Float64Array, Float64Array)> {
        let middle = self.calculate_ma(values, window)?;
        let mut upper = Vec::with_capacity(values.len());
        let mut lower = Vec::with_capacity(values.len());

        // 填充前window-1个值
        for i in 0..window-1 {
            upper.push(middle.value(i));
            lower.push(middle.value(i));
        }

        for i in window-1..values.len() {
            let window_slice: Vec<f64> = values
                .values()
                .slice(i - window + 1, window)
                .iter()
                .copied()
                .collect();

            let mean = middle.value(i);
            let variance = window_slice
                .iter()
                .map(|&x| (x - mean).powi(2))
                .sum::<f64>() / (window - 1) as f64;
            let std_dev = variance.sqrt();

            upper.push(mean + num_std * std_dev);
            lower.push(mean - num_std * std_dev);
        }

        Ok((
            Float64Array::from(upper),
            middle,
            Float64Array::from(lower),
        ))
    }

    fn calculate_rsi(
        &self,
        values: &Float64Array,
        window: usize,
    ) -> Result<Float64Array> {
        let mut rsi = Vec::with_capacity(values.len());
        let mut gains = Vec::new();
        let mut losses = Vec::new();

        // 计算价格变化
        for i in 1..values.len() {
            let change = values.value(i) - values.value(i - 1);
            gains.push(change.max(0.0));
            losses.push((-change).max(0.0));
        }

        // 填充第一个值
        rsi.push(50.0);

        // 计算初始平均收益和损失
        let mut avg_gain = gains.iter().take(window).sum::<f64>() / window as f64;
        let mut avg_loss = losses.iter().take(window).sum::<f64>() / window as f64;

        // 计算RSI
        for i in window..values.len() - 1 {
            // 更新平均收益和损失
            avg_gain = (avg_gain * (window - 1) as f64 + gains[i]) / window as f64;
            avg_loss = (avg_loss * (window - 1) as f64 + losses[i]) / window as f64;

            // 计算相对强度
            let rs = if avg_loss == 0.0 {
                100.0
            } else {
                avg_gain / avg_loss
            };

            // 计算RSI
            let rsi_value = 100.0 - (100.0 / (1.0 + rs));
            rsi.push(rsi_value);
        }

        Ok(Float64Array::from(rsi))
    }

    fn calculate_atr(
        &self,
        high: &Float64Array,
        low: &Float64Array,
        close: &Float64Array,
        window: usize,
    ) -> Result<Float64Array> {
        let mut tr = Vec::with_capacity(high.len());
        let mut atr = Vec::with_capacity(high.len());

        // 计算真实范围（TR）
        tr.push(high.value(0) - low.value(0)); // 第一个TR只使用当日高低差

        for i in 1..high.len() {
            let h_l = high.value(i) - low.value(i);
            let h_pc = (high.value(i) - close.value(i - 1)).abs();
            let l_pc = (low.value(i) - close.value(i - 1)).abs();
            tr.push(h_l.max(h_pc).max(l_pc));
        }

        // 计算ATR
        // 第一个ATR使用简单平均
        let first_atr = tr.iter().take(window).sum::<f64>() / window as f64;
        atr.push(first_atr);

        // 计算剩余的ATR值
        for i in window..tr.len() {
            let previous_atr = atr[i - 1];
            let current_tr = tr[i];
            let new_atr = (previous_atr * (window - 1) as f64 + current_tr) / window as f64;
            atr.push(new_atr);
        }

        Ok(Float64Array::from(atr))
    }
}